#!/usr/bin/env python

# Author: Edward Roualdes
#         [2012.12.13]
#         University of Kentucky

import utils
import sys
import os.path
from collections import defaultdict

def num2let(x):
    """change numbers to letters for appropriate STRUCTURE input."""
    # check for capital letters, since I don't know
    # if they will make it this far
    if x == 'a' or x == 'A':
        return str(1)
    elif x == 'c' or x == 'C':
        return str(2)
    elif x == 'g' or x == 'G':
        return str(3)
    elif x == 't' or x == 'T':
        return str(4)
    elif x == '-' or x == 'N' or x == 'n': 
        return str(-9)      
    else:
        print 'an unexpected letter!', x
        sys.exit()


def mkStruct(outdir):
    """make input file for STRUCTURE based on the NextAllele aligned haplotypes."""
    print "Making STRUCTURE file..."
    allbases = ['A', 'C', 'G', 'T', 'a', 'c', 'g', 't', '-']
    bases = ['A', 'C', 'G', 'T', 'a', 'c', 'g', 't']
    snps = defaultdict(list)
    position1 = []
    position2 = []
    # dat = utils.files2dict(os.path.join(outdir, 'haplotypes'), ".*\\.fasta$")
    # dat = utils.files2dict(os.path.join(outdir, 'haplotypes'), ".*\\.fna$")
    dat = utils.files2dict(os.path.join(outdir, 'haplotypes'), ".*\\.aligned$")

    # make data into list form; er: not sure this is optimal
    # but until I'm comfortable with fn, it stays
    for data, ind, loc in dat:
        for k, v in data.iteritems():
            data[k] = list(v)
        
        data_keys = data.keys() #list keys for data
        snps_keys = snps.keys() #list keys for snps
        if (len(snps_keys) == 0):
            num_snps_sofar=0
        else:
            num_snps_sofar=len(snps[snps_keys[0]])
                
        # go through all positions and find (record) snps
        seqlength = len(data[data_keys[0]])
        for i in range(seqlength):
            snpbases=[]
            snpbases_all=[]

            for k, v in data.iteritems(): #put all bases for that position in array
                if len(v) > i:
                    if v[i] in bases:
                        snpbases.append(v[i])
                    if v[i] in allbases: #put all bases + indels for that position in array
                        snpbases_all.append(v[i])
                        
            if (len(set(snpbases)) > 1): #if polymorphic position
                position1.append(loc) #record alignment (alignment) and position info (i)
                position2.append(i+1)
                for k, v in data.iteritems():
                    if snps[k]:                         #if the individual / allele is already there
                        if len(v) > i:
                            snps[k].append(v[i])                #add snp to snps
                        else:
                            snps[k].append('n')        #now add snp to snps				
                            
                    else:
                        for j in range(num_snps_sofar):   #add indiv / allele to snps dictionary with n's
                            snps[k].append('n')
                        snps[k].append(v[i])        #now add snp to snps
                            
                    
                            
        # even up number of snps for each indiv / allele
        max_length=0
        for k, v in snps.iteritems():       #get most number of snps
            if len(v)>max_length:
                max_length = len(v)
                        
        for k, v in snps.iteritems():       #add n's to any set of snps that isn't long enough due to indiv not being in alignment
            while len(v)<max_length:
                snps[k].append('n')

    
    # calculations finished, now write STRUCTURE file
    fileout = open(os.path.join(outdir, 'snps.txt'), 'w')
    seqlength = len(snps[snps_keys[0]])
    fileout.write("%s\t" %"individual")

    for i in range(seqlength):
        fileout.write("%s:%s\t" %(position1[i], str(position2[i]))) # write column titles
    fileout.write("\n")
    
    for indiv in sorted(snps.keys()):
        fileout.write("%s\t" %indiv)
        sequence = map(num2let, snps[indiv]) # convert letters to numbers
        for i in sequence:
            fileout.write("%s\t" % i)
        fileout.write("\n")

    return

if __name__ == "__main__":
    import sys

    outdir = '/Users/easy-e/weisrock/nextallele/misc/test/'
    mkStruct(outdir)


